package demo;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.CtNewMethod;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Objects;

/**
 * @author
 * @date 2021年12月13日 19:40
 * @description
 */
public class MyTransformer implements ClassFileTransformer {


    final static String prefix = "\nlong startTime = System.currentTimeMillis();\n";
    final static String postfix = "\nlong endTime = System.currentTimeMillis();\n";


    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (Objects.isNull(className)) {
            return null;
        }

        if (className.startsWith("java") || className.startsWith("sun")
                || className.contains("jetbrains")
                || className.contains("intellij")) {
            return null;
        }
        className = className.replace("/", ".");

        CtClass ctClass = null;

        try {

            ctClass = ClassPool.getDefault().get(className); // 使用全称,用于取得字节码类<使用javassist>
            for (CtMethod method : ctClass.getDeclaredMethods()) {
                String methodName = method.getName();
                String newMethodName = methodName + "$old";
                method.setName(newMethodName);

                CtMethod newMethod = CtNewMethod.copy(method, methodName, ctClass, null);
                // 构建新的方法体
                StringBuilder bodyStr = new StringBuilder();
                bodyStr.append("{");
                bodyStr.append("System.out.println(\"==============Enter Method: " + className + "." + methodName + " ==============\");");
                bodyStr.append(prefix);
                bodyStr.append(newMethodName + "($$);\n");// 调用原有代码，类似于method();($$)表示所有的参数
                bodyStr.append(postfix);
                bodyStr.append("System.out.println(\"==============Exit Method: " + className + "." + methodName + " Cost:\" +(endTime - startTime) +\"ms " + "===\");");
                bodyStr.append("}");

                newMethod.setBody(bodyStr.toString());// 替换新方法
                ctClass.addMethod(newMethod);// 增加新方法

            }
            return ctClass.toBytecode();
        } catch (Exception e) {
            System.out.println("出错类名:" + className);
            e.printStackTrace();
        }
        return null;
    }
}
